from typing import List

import imagehash
import wx
from PIL import Image

from event_bus import EventBus
from models import Sample
from uis import SampleEdit
from view_models import project_state, SourceCollect, SourceUnit, SampleUnit
from wxutils import call_after
from .image_mark import ImageMarkPanel


class SampleManage(SampleEdit):
    _current_page: int = 1
    _total_pages: int = 1
    _current_sources: SourceCollect
    _current_state = None

    def __init__(self, parent=None):
        super().__init__(parent)

        mark_container = self.persist_sample.GetContainingSizer()
        self.marker = ImageMarkPanel(self)
        self.marker.enable_amming(True)
        mark_container.Add(self.marker, wx.GBPosition(1, 0), wx.GBSpan(1, 3), wx.ALL | wx.EXPAND, 5)
        mark_container.AddGrowableCol(2)
        mark_container.AddGrowableRow(1)

        self.switch_enables()

        self.Bind(wx.EVT_TREE_SEL_CHANGED, self.on_source_select, self.source_list)
        self.Bind(wx.EVT_BUTTON, self.on_persist, self.persist_sample)
        self.Bind(wx.EVT_BUTTON, self.on_delete, self.remove_sample)
        self.Bind(wx.EVT_COMBOBOX, self.on_add_mark, self.mark)
        self.mark.Bind(wx.EVT_TEXT_ENTER, self.on_add_mark)
        self.mark_list.Bind(wx.EVT_KEY_UP, self.on_del_mark)
        self.Bind(wx.EVT_BUTTON, self.on_select_all, self.sample_all)
        self.Bind(wx.EVT_BUTTON, self.prev_page, self.prev_slice)
        self.Bind(wx.EVT_BUTTON, self.next_page, self.next_slice)
        self.Bind(wx.EVT_BUTTON, self.first_page, self.top_slice)
        self.Bind(wx.EVT_BUTTON, self.last_page, self.tail_slice)
        self.Bind(wx.EVT_BUTTON, self.jump_page, self.page_jump)

        EventBus().register_topic(topic="sample_source_updated", handler=self.load_sources)
        EventBus().register_topic(topic="categories_updated", handler=self.update_categories)
        EventBus().register_topic(topic="project_changed", handler=self.when_project_changed)

    @call_after
    def when_project_changed(self, _):
        EventBus().post_topic(event=dict(change_project=True), topic="sample_source_updated")
        EventBus().post_topic(event=dict(change_project=True), topic="categories_updated")

    def switch_enables(self, add_mark=False, del_mark=False, add_sample=False, del_sample=False):
        self.mark.Enable(add_mark)
        self.mark_list.Enable(del_mark)
        self.persist_sample.Enable(add_sample)
        self.remove_sample.Enable(del_sample)

    def load_sources(self, _):
        self._current_state = project_state.load_sources()
        self.load_paged_sources()

    def load_paged_sources(self, page=None, previous=None):
        self._current_sources, self._current_page, self._total_pages = self._current_state.send(page)
        self.source_list.DeleteAllItems()
        root = self.source_list.AddRoot("样本源")
        self.slice_summary.SetLabel(f"共{self._total_pages}页")
        self.slice_page.SetMax(int(self._total_pages))
        self.slice_page.SetValue(int(self._current_page))
        current_sources = self._current_sources.walk()
        try:
            item, parent = current_sources.send(None)
            while True:
                if parent is None:
                    item, parent = current_sources.send(self.source_list.AppendItem(root, item.filename))
                else:
                    item, parent = current_sources.send(self.source_list.AppendItem(parent, item.fingerprint))
        except StopIteration:
            pass
        self.source_list.ExpandAll()
        if previous is not None:
            previous_selected = self._current_sources[previous]
            if previous_selected is not None and isinstance(previous_selected, SourceUnit):
                self.source_list.SelectItem(previous_selected.item)
            elif previous_selected is not None and isinstance(previous_selected, SampleUnit):
                self.source_list.SelectItem(previous_selected.parent)
            else:
                self.marker.clear_image()
                self.update_mark_list([])
                self.switch_enables()
        else:
            self.marker.clear_image()
            self.update_mark_list([])
            self.switch_enables()

    def on_source_select(self, event: wx.TreeEvent):
        selected = self._current_sources[event.GetItem()]
        if not hasattr(selected, "parent") and selected is not None:
            self.marker.load_image(selected.filepath)
            self.update_mark_list([])
            self.switch_enables(add_mark=True, del_mark=True, add_sample=True)
        elif hasattr(selected, "parent") and selected is not None:
            parent = self._current_sources[selected.parent]
            self.marker.load_image(parent.filepath, crop=(selected.crop_start[0],
                                                          selected.crop_start[1],
                                                          selected.crop_size[0],
                                                          selected.crop_size[1]))
            self.update_mark_list([c for c in selected.categories])
            self.switch_enables(add_mark=True, del_mark=True, add_sample=True, del_sample=True)
        else:
            self.marker.clear_image()
            self.update_mark_list([])
            self.switch_enables()

    def update_categories(self, _):
        project_state.refresh_categories()
        self.refresh_page()
        self.update_mark_list()

    def update_mark_list(self, checked: List[str] = None):
        categories = {item for item in project_state.categories}
        previous_checked = set(self.mark_list.GetCheckedStrings()) & categories
        self.mark_list.SetItems(sorted(categories))
        if checked is None:
            self.mark_list.SetCheckedStrings(list(previous_checked))
        else:
            self.mark_list.SetCheckedStrings(checked)

    def on_persist(self, _):
        selected = self._current_sources[self.source_list.GetSelection()]
        if selected is None:
            wx.MessageBox("尚未选择样本源图片。", "未选择源")
            return
        if len(self.mark_list.GetCheckedItems()) == 0:
            wx.MessageBox("尚未对样本图像进行分类。", "未分类")
            return
        crop_x, crop_y, crop_width, crop_height = self.marker.crop_info
        if crop_width == 0 or crop_height == 0:
            wx.MessageBox("尚未选择样本源裁切范围。", "未裁切源")
            return
        if isinstance(selected, SampleUnit):
            Sample.delete().where(Sample.target_fingerprint == selected.fingerprint).execute()
            selected = self._current_sources[selected.parent]
        source_image = Image.open(selected.filepath)
        cropped_image = source_image.crop((max([0, crop_x]),
                                           max([0, crop_y]),
                                           min([crop_x + crop_width, source_image.width]),
                                           min([crop_y + crop_height, source_image.height])))
        new_image = Image.new(mode="RGBA", size=(crop_width, crop_height))
        dest_x = abs(crop_x) if crop_x < 0 else 0
        dest_y = abs(crop_y) if crop_y < 0 else 0
        new_image.alpha_composite(cropped_image.convert("RGBA"), (dest_x, dest_y))
        image_hash = imagehash.dhash(new_image)
        if Sample.select().where(Sample.target_fingerprint == str(image_hash)).count() > 1:
            wx.MessageBox("样本图像指纹已存在，存在潜在样本重复。", "样本可能重复")
            return
        Sample.create(source=selected.filepath,
                      source_fingerprint=selected.fingerprint,
                      crop_x=crop_x,
                      crop_y=crop_y,
                      crop_width=crop_width,
                      crop_height=crop_height,
                      target_fingerprint=str(image_hash),
                      categories=self.mark_list.GetCheckedStrings())
        self.load_paged_sources(page=self._current_page, previous=selected.fingerprint)
        EventBus().post_topic(topic="categories_updated", event=None)
        EventBus().post_topic(topic="sample_updated", event=None)

    def on_delete(self, _):
        selected = self._current_sources[self.source_list.GetSelection()]
        if isinstance(selected, SampleUnit):
            sample = Sample.get(Sample.target_fingerprint == selected.fingerprint)
            parent = sample.source_fingerprint
            Sample.delete().where(Sample.target_fingerprint == selected.fingerprint).execute()
            self.load_paged_sources(page=self._current_page, previous=parent)
        EventBus().post_topic(topic="categories_updated", event=None)
        EventBus().post_topic(topic="sample_updated", event=None)

    def on_add_mark(self, _):
        pending = self.mark.GetValue().strip()
        if len(pending) > 0 and pending not in self.mark_list.GetItems():
            self.mark_list.Append(pending)
            checked = list(self.mark_list.GetCheckedStrings())
            checked.append(pending)
            self.mark_list.SetCheckedStrings(checked)
            self.mark.SetValue("")

    def on_del_mark(self, event: wx.KeyEvent):
        if event.GetKeyCode() in [wx.WXK_DELETE, wx.WXK_BACK]:
            selected_category = self.mark_list.GetStringSelection()
            cates = Sample.categories.children().alias("cates")
            exists = Sample.select(Sample.target_fingerprint) \
                .from_(Sample, cates) \
                .where(cates.c.value == selected_category) \
                .count()
            if exists > 0:
                wx.MessageBox("要删除的特征已在其他样本中使用，请先清理样本特征后再进行删除操作。", "特征已使用")
                return
            project_state.current.categories.remove(selected_category)
            EventBus().post_topic(None, topic="categories_updated")

    def on_select_all(self, _):
        self.marker.select_all()

    def prev_page(self, _):
        self.load_paged_sources(self._current_page - 1)

    def next_page(self, _):
        self.load_paged_sources(self._current_page + 1)

    def first_page(self, _):
        self.load_paged_sources(1)

    def last_page(self, _):
        self.load_paged_sources(self._total_pages)

    def jump_page(self, _):
        self.load_paged_sources(int(self.slice_page.GetValue()))

    def refresh_page(self):
        selected = self._current_sources[self.source_list.GetSelection()]
        if selected is not None:
            self.load_paged_sources(page=self._current_page, previous=selected.fingerprint)
        else:
            self.load_paged_sources(page=self._current_page)
